import pandas as pd
import numpy as np
from sklearn.svm import SVR
from sklearn.metrics import mean_squared_error, mean_absolute_error
from sklearn.preprocessing import StandardScaler
import matplotlib.pyplot as plt

df = pd.read_csv('./dataset/pwv/pwv.csv')
df['utc_time'] = pd.to_datetime(df['utc_time'])

window_size = 36
forecast_horizon = 36  # 未来3小时

X_list = []
y_list = []

for i in range(window_size, len(df) - forecast_horizon):
    past_pwv = df['pwv'].iloc[i - window_size:i].values

    # 1小时内最小值和最大值
    pwv_1h = past_pwv[-12:]
    pwv_min_1h = np.min(pwv_1h)
    pwv_max_1h = np.max(pwv_1h)
    pwv_max_diff_1h = pwv_max_1h - pwv_min_1h

    # 1小时内最大斜率
    pwv_diff_1h = np.diff(pwv_1h)
    pwv_slopes = pwv_diff_1h
    pwv_rate_max_1h = np.max(np.abs(pwv_slopes))  # 最大斜率绝对值

    # 历史平均降雨
    past_rain_mean = df['rain'].iloc[i - window_size:i].mean()

    # 未来3小时的平均降雨量
    future_rain_avg = df['rain'].iloc[i:i + forecast_horizon].mean()

    # 组合特征
    features = np.concatenate([past_pwv, [pwv_max_diff_1h, pwv_rate_max_1h, past_rain_mean]])

    X_list.append(features)
    y_list.append(future_rain_avg)

X = np.array(X_list)
y = np.array(y_list)

split_idx = int(len(X) * 0.8)
X_train, X_test = X[:split_idx], X[split_idx:]
y_train, y_test = y[:split_idx], y[split_idx:]

scaler_X = StandardScaler()
X_train_scaled = scaler_X.fit_transform(X_train)
X_test_scaled = scaler_X.transform(X_test)

model = SVR(kernel='rbf', C=3, gamma='scale')
model.fit(X_train_scaled, y_train)
y_pred = model.predict(X_test_scaled)
y_pred = np.maximum(y_pred, 0)  # 小于0的置为0

mse = mean_squared_error(y_test, y_pred)
mae = mean_absolute_error(y_test, y_pred)
print(f'MSE: {mse:.3f}, MAE: {mae:.3f}')

plt.figure(figsize=(14, 5))
plt.plot(y_test, label='Actual', alpha=1)
plt.plot(y_pred, label='Predicted', alpha=1)
plt.xlabel('Sample')
plt.ylabel('Average Rainfall (3h)')
plt.title(f'3h Average Rainfall Forecast (MSE: {mse:.3f}, MAE: {mae:.3f})')
plt.legend()
plt.grid(True)
plt.show()
